import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import LabelEncoder

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
sns.set(font="SimHei", font_scale=1.2)


def load_data(file_path):
    """加载数据并返回目标列和特征列"""
    # 读取数据
    data = pd.read_csv(file_path)
    print("数据基本信息：")
    print(f"数据集形状: {data.shape}")
    print("\n数据集前5行：")
    print(data.head())

    # 分离目标列（第一列）和特征列（其他列）
    target_col = data.columns[0]
    target = data[target_col]
    features = data.drop(columns=[target_col])

    print(f"\n目标列: {target_col}")
    print(f"特征列数量: {features.shape[1]}")

    return data, target, features, target_col


def identify_feature_types(features, unique_threshold=10):
    """区分分类特征和连续特征"""
    # 初始化特征类型列表
    categorical_features = []
    continuous_features = []

    for col in features.columns:
        # 检查数据类型
        if features[col].dtype == 'object':
            categorical_features.append(col)
        else:
            # 对于数值型特征，通过唯一值数量判断是否为分类特征
            unique_count = features[col].nunique()
            if unique_count <= unique_threshold:
                categorical_features.append(col)
            else:
                continuous_features.append(col)

    print(f"\n分类特征数量: {len(categorical_features)}")
    print(f"连续特征数量: {len(continuous_features)}")
    print(f"分类特征: {categorical_features}")
    print(f"连续特征: {continuous_features}")

    return categorical_features, continuous_features


def analyze_categorical_features(data, target_col, categorical_features):
    """分析分类特征对目标的影响并可视化"""
    print("\n开始分析分类特征...")

    # 创建一个大画布
    plt.figure(figsize=(20, 5 * len(categorical_features)))

    for i, feature in enumerate(categorical_features, 1):
        plt.subplot(len(categorical_features), 2, 2 * i - 1)
        # 绘制特征分布
        sns.countplot(x=feature, data=data)
        plt.title(f'{feature} 分布')
        plt.xticks(rotation=45)

        plt.subplot(len(categorical_features), 2, 2 * i)
        # 绘制特征与目标的关系
        if pd.api.types.is_numeric_dtype(data[target_col]):
            # 目标是数值型，使用箱线图
            sns.boxplot(x=feature, y=target_col, data=data)
            plt.title(f'{feature} 与 {target_col} 的关系')
        else:
            # 目标是分类型，使用堆叠条形图
            crosstab = pd.crosstab(data[feature], data[target_col], normalize='index')
            crosstab.plot(kind='bar', stacked=True, ax=plt.gca())
            plt.title(f'{feature} 与 {target_col} 的关系')
        plt.xticks(rotation=45)

    plt.tight_layout()
    plt.savefig('categorical_features_analysis.png', dpi=300)
    plt.show()

    # 计算分类特征与目标的相关性（需要先编码）
    print("\n分类特征与目标的相关性：")
    encoded_data = data.copy()
    le = LabelEncoder()
    for col in [target_col] + categorical_features:
        encoded_data[col] = le.fit_transform(encoded_data[col])

    corr_matrix = encoded_data[categorical_features + [target_col]].corr()
    print(corr_matrix[target_col].sort_values(ascending=False))

    # 绘制相关性热图
    plt.figure(figsize=(20, 16))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f')
    plt.title('分类特征与目标的相关性热图')
    plt.tight_layout()
    plt.savefig('categorical_correlation.png', dpi=300)
    plt.show()


def analyze_continuous_features(data, target_col, continuous_features):
    """分析连续特征对目标的影响并可视化"""
    print("\n开始分析连续特征...")

    # 显示连续特征的基本统计信息
    print("\n连续特征的基本统计信息：")
    print(data[continuous_features].describe())

    # 创建一个大画布
    plt.figure(figsize=(20, 6 * len(continuous_features)))

    for i, feature in enumerate(continuous_features, 1):
        # 绘制直方图
        plt.subplot(len(continuous_features), 3, 3 * i - 2)
        sns.histplot(data[feature], kde=True)
        plt.title(f'{feature} 分布')

        # 绘制与目标的散点图
        plt.subplot(len(continuous_features), 3, 3 * i - 1)
        if pd.api.types.is_numeric_dtype(data[target_col]):
            sns.scatterplot(x=feature, y=target_col, data=data, alpha=0.5)
            # 添加回归线
            sns.regplot(x=feature, y=target_col, data=data, scatter=False, color='red')
        else:
            sns.violinplot(x=target_col, y=feature, data=data)
        plt.title(f'{feature} 与 {target_col} 的关系')

        # 绘制箱线图（查看异常值）
        plt.subplot(len(continuous_features), 3, 3 * i)
        sns.boxplot(y=feature, data=data)
        plt.title(f'{feature} 箱线图（异常值检测）')

    plt.tight_layout()
    plt.savefig('continuous_features_analysis.png', dpi=300)
    plt.show()

    # 计算连续特征与目标的相关性
    if pd.api.types.is_numeric_dtype(data[target_col]):
        print("\n连续特征与目标的相关性：")
        corr_matrix = data[continuous_features + [target_col]].corr()
        print(corr_matrix[target_col].sort_values(ascending=False))

        # 绘制相关性热图
        plt.figure(figsize=(10, 8))
        sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f')
        plt.title('连续特征与目标的相关性热图')
        plt.tight_layout()
        plt.savefig('continuous_correlation.png', dpi=300)
        plt.show()
    else:
        print("\n目标变量为分类型，已通过小提琴图展示与连续特征的关系")


def main(file_path):
    """主函数：执行特征工程分析流程"""
    # 加载数据
    data, target, features, target_col = load_data(file_path)

    # 识别特征类型
    categorical_features, continuous_features = identify_feature_types(features)

    # 分析分类特征
    if categorical_features:
        analyze_categorical_features(data, target_col, categorical_features)
    else:
        print("\n没有分类特征需要分析")

    # 分析连续特征
    if continuous_features:
        analyze_continuous_features(data, target_col, continuous_features)
    else:
        print("\n没有连续特征需要分析")

    print("\n特征工程分析完成！图表已保存为PNG文件。")


if __name__ == "__main__":
    # 请替换为你的数据文件路径
    data_file_path = r'../data/train.csv'  # 用户需要替换为实际的数据文件路径
    main(data_file_path)